import json
from pathlib import Path
from agents.ppo_discrete import PPODiscrete


def generate_config():
    config = {
        'use_GPU': True,
        'max_steps': 4e5,
        'eval_freq': 5000,
        'num_workers': 1,
        'batch_size': 2048,
        'mini_batch_size': 64,
        'epochs': 5,
        'actor_lr': 3e-4,
        'critic_lr': 3e-4,
        'gamma': 0.99,
        'lamda': 0.95,
        'clip_epsilon': 0.2,
        'grad_clip': 0.5,
        'entropy_coef': 0.01,  # 0.01
        'hidden_dim': 64,
        'use_orthogonal_init': True,
        'use_tanh': True,
        'use_state_norm': True,
        'use_adv_norm': True,
        'use_lr_decay': True,
    }
    return config



if __name__ == "__main__":
    env_list = ['CartPole-v1', 'Acrobot-v1', 'MountainCar-v0', 'LunarLander-v3']
    env = env_list[3]
    seed_list = [0, 10]
    config = generate_config()
    results = []
    for seed in seed_list:
        config['env_name'] = env
        config['seed'] = seed
        print(f"ENV:{env}  AGENT:PPODiscrete  SEED:{seed}")
        agent = PPODiscrete(config)
        evaluations = agent.run()
        results.append([evaluations, agent.get_score_required_to_win(), seed])
        print('-' * 100)
        print('-' * 100)
    data = {'PPODiscrete': results}
    path = Path(f'../results/data/{env}_PPODiscrete.json')
    json_str = json.dumps(data, indent=4)
    path.write_text(json_str)
